package kafkax

import (
	"context"
	"errors"
	"github.com/Shopify/sarama"
	"github.com/zeromicro/go-zero/core/logx"
	"time"
)

func NewConsumer(addr []string, topic string, group string, f func(message *sarama.ConsumerMessage) error) (*consumer, error) {
	if len(addr) == 0 || len(topic) == 0 || f == nil {
		return nil, errors.New("params required")
	}

	return &consumer{
		addr:     addr,
		topic:    topic,
		group:    group,
		f:        f,
		stopChan: make(chan bool),
	}, nil
}

type consumer struct {
	addr     []string  // 服务端地址
	topic    string    // 主题
	group    string    // 消费分组
	stopChan chan bool // 结束
	isStop   bool
	f        func(message *sarama.ConsumerMessage) error
}

func (k *consumer) Close() {
	k.isStop = true
	close(k.stopChan)
}

func (k *consumer) Run(workerNum int) error {
	if workerNum <= 0 {
		return errors.New("worker required")
	}

	// 默认自动提交,  sess.MarkMessage消息标记后才会被提交
	config := sarama.NewConfig()
	config.Consumer.Offsets.AutoCommit.Enable = false
	config.Consumer.Offsets.Initial = sarama.OffsetOldest
	config.Consumer.Return.Errors = true

	consumerGroup, err := sarama.NewConsumerGroup(k.addr, k.group, config)
	if err != nil {
		return errors.New("NewConsumerGroup error:" + err.Error())
	}
	defer consumerGroup.Close()

	go func() {
		// 自定义的consumer f返回错误时,会走这里
		for err := range consumerGroup.Errors() {
			logx.Error("consumerGroup errors", err)
		}
	}()

	ctx := context.Background()
	topics := []string{k.topic}
	for !k.isStop {
		handler := &consumerGroupHandler{
			workerNum: workerNum,
			f:         k.f,
			stopChan:  k.stopChan,
		}

		// `Consume` should be called inside an infinite loop, when a
		// server-side rebalance happens, the consumer session will need to be
		// recreated to get the new claims
		logx.Info("group.Consume start")
		err := consumerGroup.Consume(ctx, topics, handler)
		logx.Info("group.Consume end", err)

		time.Sleep(time.Second)
	}
	logx.Info("consume run stop", err)
	return nil
}
